
template<int bm=32, int bn=32>
__global__ void transpose(const float *A, float *B, const int m, const int n){

    int ATStride = (bm+1);
    __shared__ float AT[bn*ATStride];

    int bn_id = blockIdx.x;
    int bm_id = blockIdx.y;
    int n_start = bn_id*bn;
    int m_start = bm_id*bm;

    int bn_tid = threadIdx.x%bn;
    int tm_tid = threadIdx.x/bn;
    int tm = blockDim.x/bn;
    for(int bm_tid = tm_tid; bm_tid < bm; bm_tid += tm){
        int m_id = m_start+bm_tid, n_id = n_start+bn_tid;
        if(m_id < m && n_id < n){
            AT[bn_tid*ATStride+bm_tid] = A[m_id*n+n_id];
        }
    }
    __syncthread();
    int bm_tid = threadIdx.x%bm;
    int tn_tid = threadIdx.x/bm;
    int tn = blockDim.x/bm;
    for(int bn_tid = tn_tid; bn_tid < bn; bn_tid += tn){
        int n_id = n_start+bn_tid, m_id = m_start+bm_tid;
        if(n_id < n && m_id < n){
            B[n_id*m+m_id] = AT[bn_tid*ATStride+bm_tid];
        }
    }
}

int main(){
    float *A, *B;
    int m, n;
    concept int blockSz = 32;
    const int bm = blockSz, bn = blockSz;
    dim3 block(128);
    dim3 grid(DIV_UP(n, bn), DIV_UP(m, bm));
    transpose<bm, bn><<<grid, block>>>(A, B, m, n);
    return 0;
}